from __future__ import absolute_import, division, print_function

import base64
import os
import tarfile
import tempfile

import matplotlib.pyplot as plt
import numpy as np
import PIL
import tensorflow as tf
import tensorflow.contrib.slim as slim
import tensorflow.contrib.slim.nets as nets
from six.moves import urllib
from six.moves.urllib.request import urlretrieve

tf.reset_default_graph()
tf.logging.set_verbosity(tf.logging.ERROR)
sess = tf.InteractiveSession()

image = tf.Variable(tf.zeros((299, 299, 3)))

def network(image, reuse):
    preprocessed = tf.mul(tf.sub(tf.expand_dims(image, 0), 0.5), 2.0)
    arg_scope = nets.inception.inception_v3_arg_scope(weight_decay=0.0)
    with slim.arg_scope(arg_scope):
        logits, end_points = nets.inception.inception_v3(preprocessed, 1001, is_training=False, reuse=reuse)
        logits = logits[:, 1:]  # ignore background class
        probs = tf.nn.softmax(logits)  # probabilities
    return logits, probs, end_points

logits, probs, end_points = network(image, reuse=False)

# data_dir = tempfile.mkdtemp()
data_dir = '.'
checkpoint_filename = os.path.join(data_dir, 'inception_v3.ckpt')
if not os.path.exists(checkpoint_filename):
    inception_tarball, _ = urlretrieve(
        'http://download.tensorflow.org/models/inception_v3_2016_08_28.tar.gz')
    tarfile.open(inception_tarball, 'r:gz').extractall(data_dir)

restore_vars = [
    var for var in tf.all_variables() if var.name.startswith('InceptionV3/')
]

saver = tf.train.Saver(restore_vars)
saver.restore(sess, os.path.join(data_dir, 'inception_v3.ckpt'))

def get_feature(img, feature_layer_name):
    p, feature_values = sess.run([probs, end_points], feed_dict={image: img})
    return feature_values[feature_layer_name].squeeze()

image_urls = [
    '/Users/gengy/Desktop/特征相似图/cat1.jpg',
    '/Users/gengy/Desktop/特征相似图/cat2.jpeg',
    '/Users/gengy/Desktop/特征相似图/dog1.jpg',
    '/Users/gengy/Desktop/特征相似图/dog2.jpg',
    '/Users/gengy/Desktop/特征相似图/plane1.jpeg',
    '/Users/gengy/Desktop/特征相似图/plane2.jpeg',
    '/Users/gengy/Desktop/特征相似图/plane3.jpeg'
]

plt.figure(figsize=(12, 12))
images = []
for idx, img_url in enumerate(image_urls):
    img = PIL.Image.open(img_url)
    img = img.resize((299, 299))

    plt.subplot(1, 8, idx + 1)
    plt.axis('off')
    plt.imshow(img)
    plt.title('images[{}]'.format(idx))
    images.append(img)

layer = 'PreLogits'  # 这个layer就是分类器前的最后一层，详细内容参考课程视频
features = []
for img in images:
    img = (np.asarray(img) / 255.0).astype(np.float32)
    feature = get_feature(img, layer)
    features.append(feature)
feature_vectors = np.stack(features)

score_feature_shape = 0
try:
    print(feature_vectors.shape)
    assert feature_vectors.shape == (7,2048), 'shape mismatch!'
    score_feature_shape = 10
except Exception as ex:
    print(ex)

distance_euclidean = np.sum(
    np.power(feature_vectors, 2), axis=1, keepdims=True) + np.sum(
        np.power(feature_vectors, 2), axis=1,
        keepdims=True).T - 2 * np.dot(feature_vectors, feature_vectors.T)

features_norm = feature_vectors / np.linalg.norm(
    feature_vectors, axis=1)[:, np.newaxis]
distance_cosin = np.dot(features_norm, features_norm.T)

plt.figure(figsize=(12, 21))

for idx_img, plt_img in enumerate(images):
    order_euclidean = np.argsort(distance_euclidean[idx_img])
    order_cosin = np.argsort(distance_cosin[idx_img])[::-1]

    plt.subplot(14, 8, idx_img * 8 * 2 + 1)
    plt.axis('off')
    plt.imshow(plt_img)
    
    for idx_sim, i in enumerate(order_euclidean):
        similay_img = images[i]
        plt.subplot(14, 8, idx_img * 8 * 2 + idx_sim + 2)
        plt.axis('off')
        plt.title('eucl-[{}]'.format(i))

        plt.imshow(similay_img)

    for idx_sim, i in enumerate(order_cosin):
        similay_img = images[i]
        plt.subplot(14, 8, 8 + (idx_img) * 8 * 2 + idx_sim + 2)
        plt.axis('off')
        plt.title('cos-[{}]'.format(i))

        plt.imshow(similay_img)
plt.show()

